import torch.nn as nn
import torch
import torch.nn.functional as F
class bfb(nn.Module):
    def __init__(self):
        super(bfb, self).__init__()
        self.dim=8
        self.p1=nn.Conv1d(in_channels=2, out_channels=self.dim, kernel_size=1)
        self.p2= nn.AvgPool1d(kernel_size=94 //self.dim, stride=94 //self.dim)
        self.p3=nn.Parameter(torch.diag(torch.randn(self.dim)))
        self.p4=nn.BatchNorm1d(self.dim)
        self.p5=nn.Linear(self.dim*self.dim, 2*94)
        self.p6=nn.Linear(94,3)
    def forward(self, x1,x2):
        x=torch.concat([x1,x2],dim=0)
        #print(x.shape)
        x=self.p1(x)
        x=self.p2(x)
        x=F.sigmoid(x)
        x=self.p3*x
        x=self.p4(x)
        x=x.view(-1, self.dim*self.dim)
        x = self.p5(x)
        x = x.view(-1, 2, 94)
        #x=self.p5(x.view(1, 1, -1))
        #x = self.p2(x)
        x = F.softmax(F.sigmoid(x),dim=1)
        #x = (x-torch.mean(x,dim=1))/torch.std(x,dim=1) * self.p3
        x11,x21 = torch.split(x, 1, dim=1)
        x = x11*x1+x21*x2
        '''
        if not self.training:
            print(x)
        #print(x)
        '''
        x=self.p6(x)

        x = x.view(3)
        x= F.softmax(x)
        #print(x)
        return x
def get_current_memory_gb() -> int:
    import os
    import psutil
    # 获取当前进程内存占用。
    pid = os.getpid()
    p = psutil.Process(pid)
    info = p.memory_full_info()
    return info.uss / 1024. / 1024. / 1024.
if __name__=='__main__':
    import efficet
    import docrnn
    import fitter
    c = docrnn.DOCRNN()
    d = efficet.get_efficientnet_v2('efficientnet_v2_s')
    b = fitter.fitterload(
        r'D:\old\Desktop\old\animal\barking-emotion-recognition\data\audioset_audios\0BwiOU6alvQ_0_10_cut.mp3')
    x1= d(torch.unsqueeze(torch.unsqueeze(torch.from_numpy(b[0]), dim=0), dim=0))
    x2 =c(torch.unsqueeze(torch.unsqueeze(torch.from_numpy(b[1]), dim=0), dim=0))
    b = bfb()
    x = b(x1,x2)
    print(x)
    print(get_current_memory_gb())
